#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelAsinhGrad {
public:
    __aicore__ inline KernelAsinhGrad() {}
    __aicore__ inline void Init(GM_ADDR y, GM_ADDR dy, GM_ADDR z, uint32_t totalLen, uint32_t bufferSize, uint32_t tiles, uint32_t tailLen )
    {
        //考生补充初始化代码
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->bufferSize = bufferSize;
        this->totalLen = totalLen;
        this->tiles = tiles;
        this->tailLen = tailLen;
        
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y, this->totalLen);
        dyGm.SetGlobalBuffer((__gm__ DTYPE_Y *)dy, this->totalLen);
        zGm.SetGlobalBuffer((__gm__ DTYPE_Y *)z, this->totalLen);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->bufferSize * sizeof(DTYPE_Y)*2);
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->bufferSize * sizeof(DTYPE_Y));

        pipe.InitBuffer(calcBuf, this->bufferSize*sizeof(float));
        pipe.InitBuffer(zBuf, this->bufferSize*sizeof(float));
        pipe.InitBuffer(yBuf, this->bufferSize*sizeof(float));
        LocalTensor<float> buf = calcBuf.Get<float>();
        // Sub(buf, buf, buf, bufferSize);
        // Adds(buf, buf, static_cast<float>(1), bufferSize);
        Duplicate(buf, static_cast<float>(0.5), bufferSize);

    }
    __aicore__ inline void Process()
    {
        int32_t pos = 0;
        if constexpr (std::is_same_v<DTYPE_Y, half>){
            for(int i=0;i<this->tiles;i++){
                CopyIn(pos, bufferSize);
                Compute16(bufferSize);
                CopyOut(pos, bufferSize);
                pos+=this->bufferSize;
            }
            CopyIn(pos, tailLen);
            Compute16(tailLen);
            CopyOut(pos, tailLen);
        }
        else{
            for(int i=0;i<this->tiles;i++){
                CopyIn(pos, bufferSize);
                Compute(bufferSize);
                CopyOut(pos, bufferSize);
                pos+=this->bufferSize;
            }
            CopyIn(pos, tailLen);
            Compute(tailLen);
            CopyOut(pos, tailLen);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t pos, int32_t len)
    {
        LocalTensor<DTYPE_Y> xLocal = inQueueX.AllocTensor<DTYPE_Y>();
        DataCopy(xLocal, yGm[pos], len);
        DataCopy(xLocal[bufferSize], dyGm[pos], len);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute16(int32_t len)
    {
        LocalTensor<DTYPE_Y> xLocal = inQueueX.DeQue<DTYPE_Y>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        LocalTensor<DTYPE_Y> dysrc = xLocal[bufferSize];
        LocalTensor<float> y = yBuf.Get<float>();
        LocalTensor<float> z = zBuf.Get<float>();
        LocalTensor<float> buf = calcBuf.Get<float>();

        // Cast(y, xLocal, RoundMode::CAST_NONE, len);        
        // Exp(y, y, len);
        // Div(z, buf, y, len);
        // Add(y, z, y, len);
        // Cast(yLocal, y, RoundMode::CAST_NONE, len);
        // Add(dysrc, dysrc, dysrc, len);
        // Div(yLocal, dysrc, yLocal, len);

        Cast(y, xLocal, RoundMode::CAST_NONE, len);        
        Exp(y, y, len);
        Div(z, buf, y, len);
        Axpy(z, y, static_cast<float>(0.5), len);
        Cast(yLocal, z, RoundMode::CAST_NONE, len);
        Div(yLocal, dysrc, yLocal, len);

        outQueueY.EnQue<DTYPE_Y>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void Compute(int32_t len)
    {
        LocalTensor<DTYPE_Y> xLocal = inQueueX.DeQue<DTYPE_Y>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        LocalTensor<DTYPE_Y> y = xLocal;
        LocalTensor<DTYPE_Y> dy = y[bufferSize];

        LocalTensor<DTYPE_Y> buf = calcBuf.Get<DTYPE_Y>();
        
        Muls(buf, y, static_cast<DTYPE_Y>(-1), len);
        Exp(buf, buf, len);
        Exp(y,y,len);
        Add(y, buf, y, len);
        Muls(y, y, static_cast<DTYPE_Y>(0.5), len);
        Div(yLocal, dy, y, len);

        outQueueY.EnQue<DTYPE_Y>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t pos, int32_t len)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopy(zGm[pos], yLocal, len);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    //create queue for input, in this case depth is equal to buffer num
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    //create queue for output, in this case depth is equal to buffer num
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<TPosition::VECCALC> calcBuf, zBuf, yBuf;
    GlobalTensor<DTYPE_Y> zGm,yGm,dyGm;
    uint32_t totalLen, bufferSize, tiles, tailLen;

};


extern "C" __global__ __aicore__ void asinh_grad(GM_ADDR y, GM_ADDR dy, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelAsinhGrad op;
    op.Init(y,dy,z,tiling_data.totalLen, tiling_data.bufferSize, tiling_data.tiles, tiling_data.tailLen);
    op.Process();
}